import cv2
import pytesseract
from image_processing import enhance_image, extract_paper, rotate_image_180_flip

def clean_and_join_lines(text):
    """
    清理并拼接提取的文字中的换行符。
    :param text: 提取出的文字内容
    :return: 拼接后更易读的文字内容
    """
    lines = text.split('\n')  # 按行分割
    cleaned_lines = []
    current_line = ""
    for line in lines:
        stripped_line = line.strip()  # 去除首尾空格
        if stripped_line:
            if current_line:
                current_line += stripped_line  # 拼接当前行
            else:
                current_line = stripped_line  # 开始新的一行
        else:
            if current_line:
                cleaned_lines.append(current_line)  # 保存当前行
                current_line = ""
    if current_line:
        cleaned_lines.append(current_line)  # 处理最后一行
    return "\n".join(cleaned_lines)  # 返回拼接后的文本

def identify_address(enhanced_image, extract_image, address_extract_flag):
    """
    识别地址区域中的文本。
    :param enhanced_image: 增强后的图像
    :param extract_image: 提取的图像
    :param address_extract_flag: 用于识别地址的关键字
    :return: 识别出的地址文本
    """
    data = pytesseract.image_to_data(enhanced_image, lang='chi_sim', output_type=pytesseract.Output.DICT)  # 使用Tesseract识别文字及其位置
    for i in range(len(data['text'])):
        if address_extract_flag in data['text'][i]:  # 查找包含地址关键字的文本
            (x, y, w, h) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i])  # 获取边界框
            height, width = extract_image.shape[:2]
            # 根据图片的宽高比选择抠图区域
            if height < width:
                address_area = extract_image[y - int(w / 3):y + int(h * 0.75 / 2 * 13), x - w * 3:x + w + int(w / 3)]
            else:
                address_area = extract_image[y - int(h / 2):y + h * 4, x - int(h / 3):x + int(w * 0.75 / 2 * 13)]
            cv2.imwrite("extracted_address_area.png", address_area)  # 保存抠出的地址区域

            address_text = pytesseract.image_to_string(address_area, lang='chi_sim')  # 识别地址区域的文字
            cleaned_text = clean_and_join_lines(address_text)  # 清理并拼接识别出的文本
            return cleaned_text
    return None

def extract_address_area(image_path, extracted_paper_path, address_extract_flag):
    """
    提取地址区域。
    :param image_path: 输入图像的路径
    :param extracted_paper_path: 提取后的纸张图像保存路径
    :param address_extract_flag: 用于识别地址的关键字
    :return: 识别出的地址文本
    """
    extract_paper(image_path, extracted_paper_path)  # 提取纸张
    enhanced_image = enhance_image(extracted_paper_path)  # 增强纸张
    extract_image = cv2.imread(extracted_paper_path)  # 读取增强后的纸张
    address = identify_address(enhanced_image, extract_image, address_extract_flag)  # 识别地址
    if address:
        return address
    rotate_image_180_flip(extract_image, extracted_paper_path)  # 旋转180度
    enhanced_image = enhance_image(extracted_paper_path)  # 再次增强纸张
    image = cv2.imread(extracted_paper_path)  # 读取旋转后的纸张
    address = identify_address(enhanced_image, image, address_extract_flag)  # 识别地址
    return address